#!/usr/bin/env bash

if [[ $(id -u) != 0 ]]; then
    echo ">>> Error: This script must be run as root."
    exit 1
fi

# 默认参数
NEW_SIZE="32G"
PARTITION_NUMBER="1"
NBDNODE=/dev/nbd0

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case "$1" in
        -h|--help)
            echo "Usage: $(basename $0) [options]"
            echo "Options:"
            echo "    -i IMGFILE Specify the image file to resize"
            echo "    -s SIZE    Specify the new disk size (default: ${NEW_SIZE})"
            echo "    -h, --help Show this help message and exit"
            exit 0
            ;;
        -i)
            IMGFILE="$2"
            shift 2
            ;;
        -s)
            NEW_SIZE="$2"
            shift 2
            ;;
        *)
            echo ">>> Error: Invalid option $1"
            exit 1
            ;;
    esac
done

if [[ -z "$IMGFILE" ]]; then
    echo ">>> Error: Image file not specified."
    exit 1
fi

# PARTITION_NUMBER can be update by read
function cleanup {
    echo ">>> Cleaning up..."
    qemu-nbd -d ${NBDNODE} || true
    while [[ $(lsmod | grep 'nbd.*[^0]$' 2>&1 ) != ""  ]]; do
        echo ">>> Waiting for disconnect ..."
        sleep 0.1s
    done
    lsmod | grep nbd && rmmod nbd && echo ">>> NBD module removed"
}

# 设置清理函数，确保退出时释放资源
trap cleanup EXIT

set -e

# 增大虚拟磁盘的大小
echo ">>> Resizing disk image to ${NEW_SIZE}..."
qemu-img resize ${IMGFILE} ${NEW_SIZE}

# 加载 NBD 模块
echo ">>> Loading NBD module..."
lsmod | grep nbd || modprobe nbd # max_part=8

# 列出 NBD 设备
echo ">>> Listing NBD devices..."
ls /dev/nbd*

# 连接虚拟磁盘到 NBD 设备
echo ">>> Connecting disk image to NBD device..."
qemu-nbd -c ${NBDNODE} ${IMGFILE}

while [[ ! -b ${NBDNODE}p${PARTITION_NUMBER} ]]; do
    echo ">>> Waiting for partition ${PARTITION_NUMBER} to appear..."
    sleep 0.1s
done

# 打印当前分区表
echo ">>> Current partition table:"
ls /dev/nbd0*
parted -s -- ${NBDNODE} print

read -p "Enter partition num(default: ${PARTITION_NUMBER}): " PARTITION_NUMBER

# 增大指定分区的大小，剩余2M空闲空间
echo ">>> Resizing partition ${PARTITION_NUMBER}..."
parted -s -f -- ${NBDNODE} resizepart ${PARTITION_NUMBER} -2MiB

# 调整文件系统大小以匹配新的分区大小
echo ">>> Checking filesystem on partition ${PARTITION_NUMBER}..."
e2fsck -f ${NBDNODE}p${PARTITION_NUMBER}
echo ">>> Resizing filesystem on partition ${PARTITION_NUMBER}..."
resize2fs ${NBDNODE}p${PARTITION_NUMBER}

# 再次打印分区表
echo ">>> Updated partition table:"
parted -s -- ${NBDNODE} print
